import cv2
import numpy as np
import matplotlib.pyplot as plt


def SIFT_detect(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    descriptor = cv2.xfeatures2d.SIFT_create()
    return descriptor.detectAndCompute(image, None)


def KeyPoint_match(kpsA, kpsB, featA, featB, ratio, threshold):
    matcher = cv2.DescriptorMatcher_create("BruteForce")
    raw_matches = matcher.knnMatch(featA, featB, 2)
    matches = []
    good = []
    for m in raw_matches:
        if len(m) == 2 and m[0].distance < m[1].distance * ratio:
            good.append([m[0]])
            matches.append((m[0].queryIdx, m[0].trainIdx))
    kpsA = np.float32([kp.pt for kp in kpsA])
    kpsB = np.float32([kp.pt for kp in kpsB])
    if len(matches) > 4:
        ptsA = np.float32([kpsA[i] for (i, _) in matches])
        ptsB = np.float32([kpsB[i] for (_, i) in matches])
        H, status = cv2.findHomography(ptsB, ptsA, cv2.RANSAC, threshold)
        return matches, H, good
    return None


def Stitch(imgA, imgB, kpsA, kpsB, matches, H):
    hA, wA = imgA.shape[:2]
    hB, wB = imgB.shape[:2]
    image = np.zeros((max(hA, hB), wA+wB, 3), dtype='uint8')
    image[0:hA, 0:wA] = imgB
    image = cv2.warpPerspective(image, H, (image.shape[1], image.shape[0]))
    image[0:hB, 0:wB] = imgA
    return image


if __name__ == '__main__':
    imgA = cv2.imread('2.JPG')
    imgB = cv2.imread('3.JPG')
    imgAB = np.hstack((imgA, imgB))
    cv2.imshow('Input', imgAB)

    # Feature Point Detection
    kpsA, featA = SIFT_detect(imgA)
    kpsB, featB = SIFT_detect(imgB)

    cv2.imshow('KeyPointsA', cv2.drawKeypoints(imgA, kpsA, None))
    cv2.imshow('KeyPointsB', cv2.drawKeypoints(imgB, kpsB, None))

    # Feature Point Match
    matches, H, good = KeyPoint_match(kpsA, kpsB, featA, featB, ratio=0.3, threshold=0.99)
    cv2.imshow('Match', cv2.drawMatchesKnn(imgA, kpsA, imgB, kpsB, good, None,
                                           flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS))

    # Stitching
    vis = Stitch(imgA, imgB, kpsA, kpsB, matches, H)
    cv2.imshow('Pano', vis)

    cv2.waitKey()
